#!/usr/bin/env python
#coding: utf8

'''
作业:
    基于上传的密码验证 写注册和登录 用户和密码的对应关系可以用字典存到pickle到文本上
'''

import sys
import json
import hashlib
import getpass


try:
    import cPickle as pickle
except ImportError:
    import pickle


DB_PATH = 'userdb.pkl'

def open_text(fname, mode='rt'):
    "以文本方式打开文件"
    return open(fname, mode)

def md5_encrypt(field):
    "md5字段加密函数"
    try:
        return hashlib.md5(field).hexdigest()
    except TypeError as e:
        raise e

def get_db(dbfile=DB_PATH):
    "pickle反序列化获取用户的信息"
    try:
        with open_text(dbfile) as fd:
            return pickle.load(fd)
    except IOError as e:
        raise e

class Users(object):
    "管理用户对象的类"
    def __init__(self):
        self.__db = get_db()

    def __is_exist(self, user):
        "检查用户是否存在 --> bool"
        if user in self.__db.iterkeys():
            return True
        return False

    def login(self, user, passwd):
        "用户登陆方法 --> bool"
        if not self.__is_exist(user):
            return False
        passwd = md5_encrypt(passwd)
        if passwd == self.__db.get(user, None):
            return True
        return False

    def register(self, user, passwd):
        "用户注册的方法, --> bool"
        if self.__is_exist(user):
            print("Error: user '{user}' already exists.".format(user=user))
            return False
        try:
            self.__savedb(user, passwd)
            return True
        except Exception:
            return
        
    def __savedb(self, user, passwd):
        "用户信息从内存写到文件中, pickle持久化存储"
        try:
           passwd = md5_encrypt(passwd)
           self.__db[user] = passwd
        except KeyError as e:
            raise e
    
        with open_text(DB_PATH, 'wb') as fd:
            pickle.dump(obj=self.__db, file=fd)

def login():
    "登陆入口"
    err_count = 0
    usermanager = Users()

    while err_count < 3:
        u = raw_input('Please enter your username: ')
        p = getpass.getpass('Please enter your password: ')

        result = usermanager.login(u, p)
        if result:
            print(
                '\033[32mLogin Sucess, {u}.\033[0m'.format(u=u)
            )
            break
        else:
            print(
                '\033[31mError: username or password error!\033[0m'
            )
        err_count += 1

def register():
    "注册入口"
    usermanager = Users()
    u = raw_input('Please enter your username: ')
    p = getpass.getpass('Please enter your password: ')

    result = usermanager.register(u, p)
    if result:
        print(
            '\033[32mRegister successfully. {u}\033[0m'.format(u=u)
        )
    else:
        print(
           '\033[31mRegister failed!\033[0m'
        )

def main():
    "程序入口"
    print('1. Login\n2. register\n3. exit\n')
    try:
        choose = input('please input: ')
        if choose == 1:
            login()
        elif choose == 2:
            register()
        elif choose == 3:
            sys.exit('Bye!')
        else:
            sys.exit('Invalid error!')
    except Exception:
        sys.exit('Invalid error!')

if __name__ == '__main__':
    main()